# -*- coding: utf-8 -*-
import logging
import os
import time

from openai import OpenAI

import requests
import random
import json
from hashlib import md5

from black_list import Blacklist
from scholar import extract_keywords


def fetch_translator_result_ollama(content, translator_conf):
    _api_key = translator_conf['api-key']
    _model = translator_conf['model']
    # prompt = "你现在要扮演一个专业的计算机和网络空间安全从业者，我会给你一段论文的摘要，你需要通过这个摘要提供给我对应的论文关键词。你只需要回复给我你认为最接近的三个关键词。此外，不管我给你的内容是什么语言的，你的回复都需要是英文的，且你回复的格式应当如下所示'1: keywords 1 2: keywords 2 3: keywords 3'。以下是我需要你分析的论文摘要："
    prompt = translator_conf['prompt']
    # base_url = 'http://172.17.1.51:11434/v1/chat/completions' #http://172.17.1.51:11434/v1/chat/completions
    base_url = translator_conf[
        'base_url']  # 'http://172.17.1.51:11434/v1/chat/completions' #http://172.17.1.51:11434/v1/chat/completions
    payload = {
        'model': _model,
        "messages": [
            {"role": "user", "content": f"{prompt}:{content}"},
        ]
    }
    abstract = ""
    for attempt in range(100):
        abstract = ""
        try:
            response = requests.post(base_url, json=payload)
            if response.status_code == 200:
                data = response.json()
                reply = data["choices"][0]["message"]["content"]
                abstract = extract_keywords(reply, "</think>")
                logging.info(f"Get chinese abstract success. {abstract} with reply = {reply}")
                break
            else:
                logging.warning(
                    f"Attempt {attempt + 1}: Failed to get response. Status code: {response.status_code}")
        except Exception as e:
            logging.error(f"Attempt {attempt + 1}: {e}")
        time.sleep(1)
    return abstract


def fetch_translator_result(content, translator_conf):
    if content is None or content == "not found" or content == "":
        return "not found"
    if translator_conf['preference'] == "baidu":
        return fetch_translator_result_baidu(content, translator_conf)
    elif translator_conf['preference'] == "gemini":
        return fetch_translator_result_gemini(content, translator_conf)
    elif translator_conf['preference'] == "ollama":
        return fetch_translator_result_ollama(content, translator_conf)
    else:
        logging.error("Not supported translator")
        raise Exception("Not supported translator")


def fetch_translator_result_gemini(content, translator_conf):
    _api_key = translator_conf['api-keys']
    _model = translator_conf['model']
    prompt = translator_conf['prompt']
    base_url = translator_conf['base_url']
    client = OpenAI(
        api_key=_api_key,
        base_url=base_url
    )
    response = client.chat.completions.create(
        model=_model,
        n=1,
        messages=[
            {"role": "user", "content": f"{prompt}:{content}"},
        ]
    )
    # print(response.choices[0].message.content)
    return response.choices[0].message.content


# Generate salt and sign
def make_md5(s, encoding='utf-8'):
    return md5(s.encode(encoding)).hexdigest()

translator_api_key_index = 0
def fetch_translator_result_baidu(content, translator_conf):
    appid = translator_conf['appid']
    appkey = translator_conf['appkey']
    file_name = 'baidu_translator_api_key_blacklist.json'
    # file_path = os.path.join(os.path.dirname(__file__), file_name)
    blacklist = Blacklist.load_blacklist(file_name, block_time="1M")

    # 如果黑名单文件不存在，那么就不管，如果有就必须加载并检查
    if not isinstance(appid, (str, list)) or not isinstance(appkey, (str, list)):
        raise TypeError("appid 和 appkey 必须是字符串或列表")

    if isinstance(appid, str) and isinstance(appkey, str):
    #    如果只传入了一个 appid 和 appkey，先转换成对应的list
        appid = [appid]
        appkey = [appkey]
    if len(appid) != len(appkey):
        logging.error("appid 和 appkey 列表长度不一致")
        raise ValueError("appid 和 appkey 列表长度不一致")

    id_key_pair = {}
    for i in range(len(appid)):
        id_key_pair[appid[i]] = appkey[i]

    unblock_list = blacklist.get_unblocked_items(list(id_key_pair.keys()))
    logging.info(f"baidu_translator_api_key unblock_list: {unblock_list}")

    for id in id_key_pair.keys():
        if id not in unblock_list:
        # if blacklist and id in blacklist:
            logging.error(f"appid {id} 在黑名单中，无法使用")
            id_key_pair.pop(id)

    for attempt in range(100):
        if len(id_key_pair) == 0:
            logging.error("所有的 appid 都在黑名单中，无法使用")
            raise Exception("所有的 appid 都在黑名单中，无法使用")
        appid_list = list(id_key_pair.keys())
        url = translator_conf['url']
        from_lang = 'en'
        to_lang = 'zh'
    # salt = random.randint(32768, 65536)
    # sign = make_md5(appid + content + str(salt) + appkey)
    # # Build request
    # headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    # payload = {'appid': appid, 'q': content, 'from': from_lang, 'to': to_lang, 'salt': salt, 'sign': sign}
    # # Send request
    # r = requests.post(url, params=payload, headers=headers)
    # result = r.json()
    # # Show response
    # # print(json.dumps(result, indent=4, ensure_ascii=False))
    # return result["trans_result"][0]["dst"]
    # 重试机制
        try:
            # appid = appid_list[translator_api_key_index % len(appid_list)]
            appid = blacklist.get_next_unblocked_item(appid_list)
            salt = random.randint(32768, 65536)
            appkey = id_key_pair[appid]
            sign = make_md5(appid + content + str(salt) + appkey)
            # Build request
            headers = {'Content-Type': 'application/x-www-form-urlencoded'}
            payload = {'appid': appid, 'q': content, 'from': from_lang, 'to': to_lang, 'salt': salt, 'sign': sign}
            r = requests.post(url, params=payload, headers=headers)
            if r.status_code != 200:
                logging.warning(f"请求失败，状态码：{r.status_code}，重试中（{attempt + 1}/100）")
                continue

            result = r.json()
            logging.debug(result)
            # 错误码处理
            if "error_code" in result:
                error_code = result["error_code"]
                error_msg = result.get("error_msg", "未知错误")
                logging.error(f"错误码：{error_code}，错误信息：{error_msg}")

                if error_code == "52001":  # 请求超时
                    continue
                elif error_code == "52002":  # 系统错误
                    continue
                elif error_code == "52003":  # 未授权用户
                    logging.error(f"未授权用户，请检查 appid={appid} 或服务是否开通")
                    blacklist.block_item(appid)
                    id_key_pair.pop(appid)
                    continue
                elif error_code == "54000":  # 必填参数为空
                    logging.error(f"必填参数为空，请检查参数{appid}, {payload}")
                    continue
                    # raise Exception(f"必填参数为空，请检查参数{appid}")
                elif error_code == "54001":  # 签名错误
                    logging.error(f"签名错误，请检查签名生成方法{appid}, {payload}")
                    continue
                elif error_code == "54003":  # 访问频率受限
                    logging.error(f"访问频率受限，请降低调用频率或升级服务{appid}")
                    continue
                elif error_code == "54004":  # 账户余额不足
                    logging.error(f"账户{appid}余额不足，请充值")
                    blacklist.block_item(appid)
                    id_key_pair.pop(appid)
                    continue
                elif error_code == "54005":  # 长 query 请求频繁
                    time.sleep(3)
                    continue
                elif error_code == "58000":  # 客户端 IP 非法
                    logging.error(f"客户端 IP 非法，请检查服务器 IP 配置{appid}")
                    continue
                elif error_code == "58001":  # 译文语言方向不支持
                    logging.error(f"译文语言方向不支持，请检查语言参数{appid}")
                    continue
                elif error_code == "58002":  # 服务当前已关闭
                    logging.error(f"服务当前已关闭，请开启服务{appid}")
                    blacklist.block_item(appid)
                    id_key_pair.pop(appid)
                    continue
                elif error_code == "58003":  # 此 IP 已被封禁
                    logging.error(f"此 IP 已被封禁，请勿滥用 APPID {appid}")
                    blacklist.block_item(appid)
                    id_key_pair.pop(appid)
                    continue
                elif error_code == "90107":  # 认证未通过或未生效
                    logging.error(f"认证未通过或未生效，请检查认证状态{appid}")
                    continue
                elif error_code == "20003":  # 请求内容存在安全风险
                    logging.error(f"请求内容存在安全风险，请检查文本内容{appid}")
                    continue
                else:
                    logging.error(f"未知错误码：{error_code}，错误信息：{error_msg}")
                    continue

            return result["trans_result"][0]["dst"]

        except Exception as e:
            logging.error(f"请求异常：{e}")
            if attempt == 99:
                logging.error(f"{appid}重试次数已达上限，翻译失败")

                # raise Exception("重试次数已达上限，翻译失败")
    logging.error(f"{appid}翻译失败，重试次数已用尽")
    blacklist.save()
    # raise Exception("翻译失败，重试次数已用尽")


if __name__ == '__main__':
    # 测试代码
    translator_conf = {
        #  for ollama
        "api-key": "",
        #  for gemini and ollama
        "model": "deepseek-r1:14b",
        "prompt": "你现在需要做为一个英译中的科技英语翻译官，你精通网络空间安全、密码学、科技论文写作相关的所有知识。我会使用任何语言和你对话，你只需要回复我你翻译的中文结果即可，不需要其他任何内容，你现在不允许在回复中使用任何破折号。以下是我的第一个句子：",
        "base_url": "http://bamboo.gorio.top:11434/v1/chat/completions",
    }
    logging.basicConfig(level=logging.DEBUG)
    content = "Hello, world!"
    r = fetch_translator_result_ollama(content, translator_conf)
    logging.debug(r)